import numpy as np

from beam.util_psll import get_psll_by_phase


class Plant:
    def __init__(self, position, bit_num, beam_num):
        self.position = position  # 位置（二进制），形状为 (rows, cols)
        self.__bit_num = bit_num
        self.__beam_num = beam_num
        self.cost = self.calculate_cost(self.position)  # 计算成本

    def calculate_cost(self, phase):
        return get_psll_by_phase(phase, self.__bit_num, self.__beam_num)  # 计算成本

    def update_position(self, new_position):
        self.position = new_position  # 更新位置
        self.cost = self.calculate_cost(self.position)  # 重新计算成本


class RISIvyOptimization:
    def __init__(self, bit_num, beam_num, N, max_iterations, rows, cols):
        self.bit_num = bit_num
        self.beam_num = beam_num
        self.N = N  # 种群大小
        self.max_iterations = max_iterations  # 最大迭代次数
        self.rows = rows  # 行数
        self.cols = cols  # 列数

        self.var_size = (rows, cols)  # 决策变量矩阵大小

        self.pop = []  # 初始种群数组
        self.best_costs = np.zeros(self.max_iterations)  # 最佳成本历史
        self.convergence_curve = []  # 收敛曲线


    def initialize_population_rand(self):
        for _ in range(self.N):
            position = np.random.randint(2, size=self.var_size).astype(int)  # 初始化二进制位置
            plant = Plant(position, self.bit_num, self.beam_num)  # 创建植物对象
            self.pop.append(plant)

    def initialize_population_by_phase_init(self, phase_init):
        for _ in range(self.N):
            plant = Plant(phase_init, self.bit_num, self.beam_num)  # 创建植物对象
            self.pop.append(plant)


    def crossover(self, parent1, parent2):
        # 单点交叉
        point_row = np.random.randint(1, self.rows - 1)
        point_col = np.random.randint(1, self.cols - 1)

        child1_top_left = parent1[:point_row, :point_col]
        child1_top_right = parent2[:point_row, point_col:]
        child1_bottom_left = parent2[point_row:, :point_col]
        child1_bottom_right = parent1[point_row:, point_col:]

        child1 = np.vstack((np.hstack((child1_top_left, child1_top_right)),
                            np.hstack((child1_bottom_left, child1_bottom_right))))

        child2_top_left = parent2[:point_row, :point_col]
        child2_top_right = parent1[:point_row, point_col:]
        child2_bottom_left = parent1[point_row:, :point_col]
        child2_bottom_right = parent2[point_row:, point_col:]

        child2 = np.vstack((np.hstack((child2_top_left, child2_top_right)),
                            np.hstack((child2_bottom_left, child2_bottom_right))))

        return child1, child2

    def mutation(self, individual, mutation_rate=0.01):
        # 变异操作
        mutated_individual = individual.copy()
        for i in range(individual.shape[0]):
            for j in range(individual.shape[1]):
                if np.random.rand() < mutation_rate:
                    mutated_individual[i, j] = 1 - mutated_individual[i, j]
        return mutated_individual

    def run(self, phaseBit_init, logger):

        # """初始化返回值"""
        self.best_individual = phaseBit_init
        self.best_fitness = 0
        self.best_individual_history = []
        self.best_fitness_history = []

        # 随机初始化
        self.initialize_population_rand()
        # 使用phase_init初始化
        self.initialize_population_by_phase_init(phase_init=phaseBit_init)

        for it in range(self.max_iterations):
            costs = [plant.cost for plant in self.pop]  # 获取所有植物的成本
            best_cost = min(costs)  # 最佳成本

            newpop = []

            for i in range(self.N):
                ii = i + 1 if i < self.N - 1 else 0  # 下一个索引，循环到最后则回到第一个

                beta_1 = 1 + np.random.rand() / 2  # beta值，见论文中算法1第8行

                if self.pop[i].cost < beta_1 * self.pop[0].cost:
                    new_position1, new_position2 = self.crossover(self.pop[i].position, self.pop[ii].position)

                    # 添加变异操作
                    new_position1 = self.mutation(new_position1)
                    new_position2 = self.mutation(new_position2)

                    newsol1 = Plant(new_position1, self.bit_num, self.beam_num)  # 创建新植物对象
                    newsol2 = Plant(new_position2, self.bit_num, self.beam_num)  # 创建新植物对象

                    newpop.append(newsol1)
                    newpop.append(newsol2)
                else:
                    new_position = self.mutation(self.pop[0].position)
                    newsol = Plant(new_position, self.bit_num, self.beam_num)  # 创建新植物对象
                    newpop.append(newsol)

            # 合并种群
            self.pop.extend(newpop)

            # 排序种群
            self.pop.sort(key=lambda x: x.cost)

            # 竞争排斥（删除多余成员）
            if len(self.pop) > self.N:
                self.pop = self.pop[:self.N]

            # 存储迄今为止找到的最佳解
            best_sol = self.pop[0]
            self.best_costs[it] = best_sol.cost
            self.convergence_curve.append(best_sol.cost)

            # 记录最佳适应度曲线
            if best_sol.cost < self.best_fitness:
                self.best_individual = best_sol.position.copy()
                self.best_fitness = best_sol.cost
            self.best_fitness_history.append(self.best_fitness)
            self.best_individual_history.append(self.best_individual)

            # 打印当前最优解
            # print(f"Iteration {iteration + 1}: Best Fitness = {self.best_fitness}")
            logger.info("iteration=%d: self.best_fitness=%f" % (it, self.best_fitness))
            # logger.info("---------------------------------------------------------------------------")

        destination_fitness = self.pop[0].cost  # 目标适应度
        destination_position = self.pop[0].position  # 目标位置

        # return destination_fitness, destination_position, self.convergence_curve
        return self.best_individual, self.best_fitness, self.best_fitness_history, self.best_individual_history


